package com.celan.year2023.month05.day28;

import java.util.Comparator;
import java.util.PriorityQueue;
import java.util.Queue;

/**
 * @author Celan
 * @description TODO
 */
public class Solution {
    public int kthSmallest(int[][] mat, int k) {
        int[] a = new int[]{0};
        for (int[] row : mat)
            a = kSmallestPairs(row, a, k);
        return a[k - 1];
    }

    //373. 查找和最小的 K 对数字
    private int[] kSmallestPairs(int[] nums1, int[] nums2, int k) {
        int n = nums1.length, m = nums2.length, size = 0;
        int[] ans = new int[Math.min(k, n * m)];
        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> a[0] - b[0]);

        //[sum, nums1.idx, nums2.idx]
        pq.add(new int[]{nums1[0] + nums2[0], 0, 0});
        while (!pq.isEmpty() && size < k) {
            int[] poll = pq.poll();
            int idx1 = poll[1], idx2 = poll[2];
            //数对和
            ans[size++] = nums1[idx1] + nums2[idx2];
            if (idx2 == 0 && idx1 + 1 < n)
                pq.add(new int[]{nums1[idx1 + 1] + nums2[0], idx1 + 1, 0});
            if (idx2 + 1 < m)
                pq.add(new int[]{nums1[idx1] + nums2[idx2 + 1], idx1, idx2 + 1});
        }
        return ans;
    }
}
